###################### construct graph ##########################
###### the version : complete sub-graph with super node

import matrix
import global_var
        
### twitter api agent ###
import twitter
import MySQLdb
import cPickle as pickle

class TwitterApiAgentDB:
    
    api = twitter.Api( \
        consumer_key = "S0HrvMBtFWrnccV9mXKKeQ", \
        consumer_secret = "EVkctWVhzv1CdvT4Tkezf1pA6LU8xpk0udfgdWpNAk", \
        access_token_key = "10434312-vEHrNtZYqpT7z2iPXqPqq5JEJlRsPulHLTLZEEk3K", \
        access_token_secret = "4ifpNIckQkyBLKWw6rwJ8bMuqTYZs4W6NLiIy1Yi4")
        
    
    db = MySQLdb.connect(user='root',passwd='123',db='twitter')
    cursor = db.cursor()
    
    

    def getName(self, id):
        name = None
        
        sqltext = '''select id, name from user where id = ''' + str(id)
        cnt = self.cursor.execute(sqltext)
        
        if cnt > 0 :
            t = self.cursor.fetchone()
            return t[1]
        try:
            user = self.api.GetUser(id)
            name = user.screen_name
            
        except:
            msg = "id " + str(id) + " can not get name," \
                + " maybe the id is not available."
            print >> global_var.logger, msg
            return "unknown name"
        
        user_str = pickle.dumps(user, 2)
        sqltext = '''insert into user values (%s, %s, %s )'''
        self.cursor.execute(sqltext,(id, name, user_str))
        return name

    def getId(self, name): 
        id = None
        
        sqltext = '''select id, name from user where name = ''' + '\'' +  str(name) + '\''        
        cnt = self.cursor.execute(sqltext)
        
        if cnt > 0 :
            t = self.cursor.fetchone()
            return t[0]
        try:
            user = self.api.GetUser(name)
            id = user.id
        except:
            msg = "name " + str(name) + " can not get id," \
                + " maybe the name is not available or network error!."
            print >> global_var.logger, msg
            return -1
              
        user_str = pickle.dumps(user, 2)
        sqltext = '''insert into user values (%s, %s, %s )'''
        self.cursor.execute(sqltext,(id, name, user_str))
            
        return id
    
    
    def getUser(self, id = None, name = None):
        user = None
        if isinstance(id, int) or isinstance(id, long):
            sqltext = '''select detail from user where id = ''' + '\'' +  str(id) + '\''        
            cnt = self.cursor.execute(sqltext)
            
            if cnt > 0:               
                t = self.cursor.fetchone()
                t = t[0]
                user = pickle.loads(t)
                return user
            
            try:
                user = self.api.GetUser(id)
            except:
                msg = "id " + str(id) + " can not get user," \
                    + " maybe the id is not available."
                print >> global_var.logger, msg
                return user
        
        if isinstance(name, str):
            sqltext = '''select detail from user where name = ''' + '\'' +  str(name) + '\''        
            cnt = self.cursor.execute(sqltext)
            
            if cnt > 0:               
                t = self.cursor.fetchone()
                t = t[0]
                user = pickle.loads(t)
                return user
            
            try:
                user = self.api.GetUser(name)
            except:
                msg = "name " + str(name) + " can not get user," \
                    + " maybe the name is not available."
                print >> global_var.logger, msg
                return user
        
        user_str = pickle.dumps(user, 2)
        sqltext = '''insert into user values (%s, %s, %s )'''
        self.cursor.execute(sqltext,(user.id, user.screen_name, user_str))        
        
        return user


    def isPrivate(self, id):
        user = self.getUser(id)
        return user.protected
    
    def isHub(self, id): # followers_n > 2000
        user = self.getUser(id)
        return user.followers_count > 2000
    
    def getFFT(self, id): # followers_n friends_n tweets_n
        user = self.getUser(id)
        return (user.followers_count, user.friends_count, user.statuses_count)
    
    def getFollowers(self, id):  # 10,000 is the limit !!!
        followerids = []
        if self.isPrivate(id):
            return followerids
        
        sqltext = '''select * from followers where id = ''' + str(id)
        cnt = self.cursor.execute(sqltext)

        if cnt > 0:
            t = self.cursor.fetchone()
            list = t[1]
            followerids = pickle.loads(list)
            return followerids
            
        try:
            results = self.api.GetFollowerIDs(id)
            followerids = results['ids']
            while(results['next_cursor'] != 0):
                results = self.api.GetFollowerIDs(id, results['next_cursor'])
                followerids.extend(results['ids'])            
        except:
            name = self.getName(id)
            msg =  "" + name + "\'s followers id list can not be got, " + \
                    "maybe get data from twitter error."
            print >> global_var.logger, msg
            return followerids
        
        sqltext = '''insert into followers values (%s, %s) '''
        self.cursor.execute(sqltext, (str(id), pickle.dumps(followerids, 2)))
                    
        return followerids
        

    def getFriends(self, id):  # 10,000 is the limit !!!
        friendids = []
        if self.isPrivate(id):
            return friendis       
        
        sqltext = '''select * from friends where id = ''' + str(id)
        cnt = self.cursor.execute(sqltext)

        if cnt > 0:
            t = self.cursor.fetchone()
            list = t[1]
            friendids = pickle.loads(list)
            return friendids
            
        try:
            results = self.api.GetFriendIDs(id)
            friendids = results['ids']
            while(results['next_cursor'] != 0):
                results = self.api.GetFriendIDs(id, results['next_cursor'])
                friendids.extend(results['ids'])           
        except:
            name = self.getName(id)
            msg =  "" + name + "\'s friends id list can not be got, " + \
                    "maybe the user is private. or get data from twitter error."
            print >> global_var.logger, msg
            return friendids
        
        sqltext = '''insert into friends values (%s, %s) '''
        self.cursor.execute(sqltext, (str(id), pickle.dumps(friendids, 2)))
                    
        return friendids
    
    def loadTweets(self, user_id_or_name):
        all_twts = []
        try:
            reslen = 1
            i = 1
            while reslen != 0:              
                twts = self.api.GetUserTimeline(id=user_id_or_name, page= i, count = 200, include_rts=True, include_entities=True)
                reslen = len(twts)
                i += 1               
                all_twts.extend(twts)
        except:
            user = self.getUser(user_id_or_name, user_id_or_name)
            msg =  "" + user.screen_name + "\'s tweets loading error!"
            print >> global_var.logger, msg
            return []
        
        for twt in all_twts:
            name = twt.user.screen_name
            id = twt.user.id
            sqltext = ''' insert into tweets values (%s,%s,%s,%s,%s,%s,%s,%s)'''
            self.cursor.execute(sqltext,(twt.id, \
                                         twt.text.encode('utf-8'), name, id, twt.in_reply_to_status_id, \
                                         twt.in_reply_to_screen_name, twt.in_reply_to_user_id, \
                                         pickle.dumps(twt,2)))
        
        return all_twts


######## global var
GRAPH_SUPER_NODE_ID = 0.1

### Class:Graph Node
class GraphNode:
    id = -1
    nodes_list = None # list
    tagged = False
    
    global GRAPH_SUPER_NODE_ID
    
    

    def __init__(self, id = -1, nodes_list = None, tagged = False):
        self.id = id
        self.nodes_list = nodes_list
        self.tagged = tagged

    def printNode(self, agent):
        namelist = [];
        for id in self.nodes_list:
            if id == GRAPH_SUPER_NODE_ID :
                namelist.append('super_node')
                continue
            name = agent.getName(id)
            namelist.append(name)
        
        if self.id == GRAPH_SUPER_NODE_ID:
            name = 'super_node'
        else:
            name = agent.getName(self.id)
        msg = str(name) + '\t' + str(namelist) + str(self.tagged)
        print >> global_var.logger, msg
        return msg

    def addTag(self):
        self.tagged = True

    def removeTag(self):
        self.tagged = False

    def isTagged(self):
        return self.tagged
    
    def haveNode(self, node_id):
        return node_id in self.nodes_list

    def deleteNode(self, node_id):
        if self.haveNode(node_id):
            self.nodes_list.remove(node_id)
        else:
            return False
        
        return True        

    def addNode(self, node_id, dup = False):
        if dup:
            self.nodes_list.append(node_id)
            return
        if node_id not in self.nodes_list:
            self.nodes_list.append(node_id)

### Class:Graph
from collections import deque
class Graph:
    nodes_list = None
    global GRAPH_SUPER_NODE_ID
    
    
    
    def __init__(self):
        self.nodes_list = {}
        g_node = GraphNode(GRAPH_SUPER_NODE_ID, [GRAPH_SUPER_NODE_ID])
        self.nodes_list[GRAPH_SUPER_NODE_ID] = g_node
        
    def clear(self):
        self.nodes_list.clear()
        
    def addEdge(self, source, target):
        if source in self.nodes_list:
            self.nodes_list[source].addNode(target)
        else:
            new_node = GraphNode(source, [target])
            self.nodes_list[source] = new_node

    def construct_v2(self, agent, name, max_level = 10):
        start_id = agent.getId(name)
        
        if agent.isHub(start_id):
            msg = "%s is hub!" %(name)
            print >> global_var.logger, msg
            return
        
        if agent.isPrivate(start_id):
            msg = "%s is private!" %(name)
            print >> global_var.logger, msg
            return            
        
        cache = set()
        queue = deque()
        entity = (start_id, 1)  # id: start_id, level : 1
        queue.append(entity)
        cache.add(start_id)
        
        while len(queue) > 0:
            entity = queue.popleft()
            travel_id = entity[0]
            level = entity[1]
            
            follower_ids = agent.getFollowers(travel_id)
            friend_ids = agent.getFriends(travel_id)
            
            if level < max_level:
                for id in follower_ids:
                    try:
                        user = agent.getUser(id)
                        if user == None :
                            continue
                    except:
                        continue
                    
                    self.addEdge(travel_id, id)
                    
                    if (not agent.isPrivate(id)) and (not agent.isHub(id)):
                        if id not in cache:
                            new_entity = (id, level + 1)
                            queue.append(new_entity)                        
                    
                    cache.add(id)
                    
                for id in friend_ids:
                    try:
                        user = agent.getUser(id)
                        if user == None :
                            continue
                    except:
                        continue
                    
                    self.addEdge(id, travel_id)
                    
                    if (not agent.isPrivate(id)) and (not agent.isHub(id)):
                        if id not in cache:
                            new_entity = (id, level + 1)
                            queue.append(new_entity)
                            
                    cache.add(id)  
                    
            else: # the max_level
                for id in follower_ids:
                    if id in cache:
                        self.addEdge(travel_id, id)
                for id in friend_ids:
                    if id in cache:
                        self.addEdge(id, travel_id) 
                        
        for id in cache:
            if id not in self.nodes_list:
                g_node = GraphNode(id, [])
                self.nodes_list[id] = g_node
                    
                   
        
    
    def construct(self, agent, name, max_level = 10):
        start_id = agent.getId(name)
        
        in_nodes_list = {}

        cache = set()
        queue = deque()
        entity = (start_id, 1) # id: start_id, level : 1
        queue.append(entity)
        cache.add(start_id)
        while len(queue) > 0 :
            entity = queue.popleft()
            travel_id = entity[0]
            level = entity[1]
            
            friend_ids = agent.getFriends(travel_id)
            in_nodes_list[travel_id] = friend_ids
            
            follower_ids = agent.getFollowers(travel_id)
            if level < max_level:
                node = GraphNode(travel_id, follower_ids)  
                self.nodes_list[travel_id] = node
   
                for id in follower_ids:
                    if id not in cache :
                        new_entity = (id, level + 1)
                        queue.append(new_entity)
                        cache.add(id)
            else: # the max_level
                list = []
                for id in follower_ids:
                    if id in cache:
                        list.append(id)
                    else:
                        list.append(GRAPH_SUPER_NODE_ID)
                node = GraphNode(travel_id, list)  
                self.nodes_list[travel_id] = node


        for id, list in in_nodes_list.iteritems():
            for tid in list:
                if tid not in cache:
                    self.nodes_list[GRAPH_SUPER_NODE_ID].addNode(id, True)
                else:
                    self.nodes_list[tid].addNode(id)

    
    def printGraph(self, agent, filename = ""):
        try:
            f = file(filename, 'w')
            for node in self.nodes_list.itervalues():
                msg = node.printNode(agent)
                f.write(msg)
                f.write('\n')
            f.close()
        except:
            for node in self.nodes_list.itervalues():
                node.printNode(agent)
                
    def generateNWB_v2(self, agent, filename = ""):
        try:
            sn = -1
            f = file(filename, 'w')
            n = len(self.nodes_list)
            f.write("*Nodes " + str(n) + '\n')
            f.write("id*int label*string\n")
            index = 0
            dict = {}
            
            in_degree = {}
            out_degree = {}
            
            for node in self.nodes_list:
                in_degree[node] = 0
                index = index + 1
                dict[node] = index
                f.write(str(index) + " " + "\"" + str(node) + "\"\n")
                if node == GRAPH_SUPER_NODE_ID:
                    sn = index
            
            
            f.write("*DirectedEdges\n")
            f.write("source*int target*int weight*float\n")

            for id, node in self.nodes_list.iteritems():
                source = id
                out_degree[source] = len(node.nodes_list)
                for target in node.nodes_list:
                    f.write(str(dict[source]) + " " + str(dict[target]) + " " + str(1) + "\n")
                    in_degree[target] += 1

                        
            for id, cnt in out_degree.iteritems() :
                if id == GRAPH_SUPER_NODE_ID:
                    continue
                fo_n, fr_n, t_n = agent.getFFT(id)
                d = fo_n - cnt
                if d != 0 :
                    f.write(str(dict[id]) + " " + str(sn) + " " + str(d) + "\n")
                    
            for id, cnt in in_degree.iteritems() :
                if id == GRAPH_SUPER_NODE_ID:
                    continue
                fo_n, fr_n, t_n = agent.getFFT(id)
                d = fr_n - cnt
                if d != 0 :
                    f.write(str(sn) + " " + str(dict[id]) + " " + str(d) + "\n")
                    
            f.close()
        except:
            msg = "file " + str(filename) + " can not open!"
            print >> global_var.logger, msg
            
                       
    def generateNWB(self, agent, filename = ""):
        try:
            f = file(filename, 'w')
            n = len(self.nodes_list)
            f.write("*Nodes " + str(n) + '\n')
            f.write("id*int label*string\n")
            index = 0
            dict = {}
            for node in self.nodes_list:
                index = index + 1
                dict[node] = index
                f.write(str(index) + " " + "\"" + str(node) + "\"\n")
            f.write("*DirectedEdges\n")
            f.write("source*int target*int weight*float\n")
            for id, node in self.nodes_list.iteritems():
                source = id
                dic = {}
                for tid in node.nodes_list:
                    dic[tid] = 0
                for tid in node.nodes_list:
                    dic[tid] = dic[tid] + 1
                for tid, cnt in dic.iteritems():
                    f.write(str(dict[source]) + " " + str(dict[tid]) + " " + str(cnt) + "\n")
            f.close()
        except:
            msg = "file " + str(filename) + " can not open!"
            print >> global_var.logger, msg

import re           
def loadNwbMapping(filename, inverse = False): # have not super-node
        dic = {}
        f = file(filename, 'r')
        p = re.compile('[0-9]+ \"[0-9]+')
        for line in f:
            m = p.match(line)
            if m != None:
                line = line.replace('\"','')
                res = line.split(' ')
                i = int(res[0])
                try:
                    j = long(res[1])
                    if inverse:
                        dic[j] = i
                    else:
                        dic[i] = j
                except:
                    pass
        f.close()
        return dic
    
def getIdIndex(id, filename):  
    ret = -1
    f = file(filename, 'r')
    p = re.compile('[0-9]+ \"[0-9]+')
    for line in f:
        m = p.match(line)
        if m != None:
            line = line.replace('\"','')
            res = line.split(' ')
            i = int(res[0])
            try:
                j = long(res[1])
                if j == id:
                    ret = i
            except:
                pass
    f.close()
    print >> global_var.logger, ("start index:" + str(ret))
    return ret

def addNameForNwbFile(agent, filename):
    fout_name = filename + ".name"
    fout = file(fout_name,'w')
    f = file(filename, 'r')
    p = re.compile('[0-9]+ \"[0-9]+')
    for line in f:
        m = p.match(line)
        if m != None:
            line = line.replace('\"','')
            res = line.split(' ')
            i = int(res[0])
            try:
                j = long(res[1])
                name = agent.getName(j)
                s = str(i) + ' ' + '\"' + str(name) + '\"'
                fout.write(s)
            except:
                s = str(i) + ' ' + "\"OtherNodes\""
                fout.write(s)
            fout.write('\n')
        else:
            fout.write(line)
    f.close()
    fout.close()

def nwb2gexf(agent, nwb_filename, possibility_filename):
    head = \
'''<?xml version="1.0" encoding="UTF-8"?>
<gexf xmlns="http://www.gexf.net/1.1draft" version="1.1" xmlns:viz="http://www.gexf.net/1.1draft/viz" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://www.gexf.net/1.1draft http://www.gexf.net/1.1draft/gexf.xsd">
    <meta>
        <creator>Xie Wei</creator>
        <description>Twitter Graph</description>
    </meta>
    <graph defaultedgetype="directed">
        <attributes class="node" mode="static">
            <attribute id="twt_id" title="Tweet Id" type="biginteger"></attribute>
            <attribute id="twt_name" title="Tweet Name" type="string"></attribute>
            <attribute id="p" title="Possibility" type="double"></attribute>
            <attribute id="private" title="Private" type="boolean"></attribute>
            <attribute id="hub" title="Hub" type="boolean"></attribute>
            <attribute id="followers_n" title="Number of Followers" type="integer"></attribute>
            <attribute id="friends_n" title="Number of Friends" type="integer"></attribute>
            <attribute id="tweets_n" title="Number of Tweets" type="integer"></attribute>
        </attributes>
        <nodes>'''
    gexf_file = file(nwb_filename + ".gexf" , 'w')    
    nwb_file = file(nwb_filename, 'r')
    p_file = file(possibility_filename, 'r')
    print >> gexf_file, head
    p_dic = {}
    for line in p_file:
        res = line.split(' ')
        if len(res) == 2:
            id = int(res[0])
            p = float(res[1])
            p_dic[id] = p
    p_file.close()
    
    super_node_id = -1
    
    p1 = re.compile('[0-9]+ \"[0-9]+')
    p2 = re.compile('[0-9]+ [0-9]+')
    for line in nwb_file:
        m = p1.match(line)
        if m != None:
            line2 = line.replace('\"','')
            res = line2.split(' ')
            i = int(res[0])
            if "0.1" not in res[1]:
                j = long(res[1])
                name = agent.getName(j)
                s = \
'''          <node id="%s" label="%s">
            <attvalues>
              <attvalue for="twt_id" value="%s"></attvalue>
              <attvalue for="twt_name" value="%s"></attvalue>
              <attvalue for="p" value="%s"></attvalue>
              <attvalue for="private" value="%s"></attvalue>
              <attvalue for="hub" value="%s"></attvalue>
              <attvalue for="followers_n" value="%s"></attvalue>
              <attvalue for="friends_n" value="%s"></attvalue>
              <attvalue for="tweets_n" value="%s"></attvalue>
            </attvalues>
          </node>'''
                if i in p_dic:
                    p = p_dic[i]
                else:
                    p = 0
                
                pri = agent.isPrivate(j)
                hub = agent.isHub(j)
                fo_n, fr_n, t_n = agent.getFFT(j)
                print >> gexf_file, s %(str(i),str(name),str(j),str(name),str(p),\
                                        str(pri), str(hub), str(fo_n), str(fr_n), str(t_n))
            else:
                super_node_id = i
                
                
        if line.startswith("source"):
            print >> gexf_file, "        </nodes>"
            print >> gexf_file, "        <edges>"
            
        m = p2.match(line)
        if m != None:
            s = '''            <edge source="%s" target="%s"></edge>'''
            res = line.split(' ')                
            source = int(res[0])
            target = int(res[1])
            if super_node_id != source and super_node_id != target:
                print >> gexf_file, s %(str(source), str(target))
    nwb_file.close()
    s = \
'''        </edges>
    </graph>
</gexf>'''
    print >> gexf_file, s

    gexf_file.close()
   
import time  
import sys 
if __name__ == '__main__':
    
    if len(sys.argv) <= 1:
            name = "huixue"  
            hops = 2
            p = 0.8
    else:
        name = sys.argv[1]
        hops = int(sys.argv[2])
        p = float(sys.argv[3])
    
    agent = TwitterApiAgentDB()
    #agent = TwitterApiAgent()

       
    #filestr = "/mnt/share/graph/graph_2_0.8_crimsontrauma_1300848199.42.nwb"
    #addNameForNwbFile(agent, filestr)

    
    '''
    if True:
        sys.exit()
    print 'I am not out!'
    '''
    
            
    filestr = "/mnt/share/graph/graph_" + str(hops) + "_" \
        + str(p) + "_" + str(name) + '_' + str(time.time()) + ".nwb"
        
    logfile = file(filestr + ".log", 'w')
    
    global_var.logger = logfile
    
    graph = Graph()
    graph.construct_v2(agent, name, hops)  
    #graph.printGraph(agent)        
    graph.generateNWB_v2(agent, filename = filestr)
        
    graph.clear()

    '''   
    dic = loadNwbMapping(filestr, inverse = True)
    for id in dic.iterkeys():
        agent.loadTweets(id)
    '''    
    
   
    mat = matrix.SparseMatrix()
    mat.loadFromNwbFile(filestr)
    #########
    #mat.printInfo()
    #########
    mat = mat.prepareForRW()
         
    id = agent.getId(name)
    index = getIdIndex(id, filestr)
    
    mat.randomWalk(p, index, 1e-5, True)
    list = mat.randomWalkPair(p, index, 1e-5, filestr + ".p")
    mat.clear()
    
    nwb2gexf(agent, filestr, filestr + ".p")
    
    dic = loadNwbMapping(filestr)
    
    for t in list:
        if t[0] in dic:
            id = dic[t[0]]
            name = agent.getName(id)
            print >> logfile, name, t[1]

        
    logfile.close()